/*
 * kdtree_impl.h
 *
 *  Created on: Mar 3, 2011
 *      Author: tqlong
 */

#ifndef KDTREE_IMPL_H_
#define KDTREE_IMPL_H_
#include "kdtree.h"
#include <algorithm>
#include <limits>
#include <cmath>

#define DOUBLE_INFINITY std::numeric_limits<double>::infinity()

template<typename T, typename NS>
void KDTreeNode<T,NS>::split()
{
//  std::cerr << "split n = " << this->n() << " (" << this->lo_ << " " << this->hi_ex_ << ")\n";
  if (this->n() <= 10) return;
  std::vector<int> dims(this->dim());
  for (int i = 0; i < this->dim(); i++) dims[i] = i;     // dimension list to try

  for (int i = 0; i < this->dim(); i++) {
    int d = rand()%(this->dim()-i)+i;
    int tmp = dims[i]; dims[i] = dims[d]; dims[d] = tmp; // choose a random dimension
    d = dims[i];                                         // to split data
    if (split(d)) break;                                 // if successful, done
  }
}

template<typename T, typename NS>
void KDTreeNode<T,NS>::splitRecursive()
{
//  std::cerr << "splitrec n = " << this->n() << " (" << this->lo_ << " " << this->hi_ex_ << ")\n";
  split();
  for (int i = 0; i < this->n_child(); i++)
    this->child(i)->splitRecursive();
  combineStatistics();
//  nstats_.print();
}

template<typename T, typename NS>
double KDTreeNode<T,NS>::median5(std::vector<double>& x, int i, int n)
{
  for (int ii = i; ii < i+n; ii++)
    for (int jj = ii+1; jj < i+n; jj++)
      if (x[ii]>x[jj]) {
        double tmp = x[ii]; x[ii] = x[jj]; x[jj] = tmp;
      }
  return x[i+(n-1)/2];
}

template<typename T, typename NS>
double KDTreeNode<T,NS>::median(std::vector<double>& x, int n)
{
  if (n <= 5) return median5(x, 0, n);
  int k = 0;
  for (int i = 0; i < n; i += 5) {
    int nn = 5;
    if (nn > n-i) nn = n-i;
    x[k] = median5(x, i, nn);
    k++;
  }
  return median(x, k);
}

template<typename T, typename NS>
bool KDTreeNode<T,NS>::split(int d)
{
  typedef KDTreeNode<T,NS> KN;
  std::vector<double> x(this->n());
  for (int i = 0; i < this->n(); i++)                    // extract dimention d
    x[i] = this->point(i).get(d);                        // value and sort them

//  std::sort(x.begin(), x.end());
//  double val = x[(x.size()-1)/2];                            // take median value
  double val = median(x, x.size());
  int k = 0;                                             // and try to split the node
//  std::cerr << "  val = " << val << "\n";
  for (int i = 0; i < this->n(); i++)
    if (this->point(i).get(d) <= val) {
      this->swapIndex(k, i);
      k++;
    }
  if (k >= this->n()) return false;                      // if left node has all points, failure

  KN *left = new KN(*this, 0, k);                        // otherwise, create left and right nodes
  KN *right = new KN(*this, k, this->n());
  this->addChild(*left);
  this->addChild(*right);
  return true;
}

template<typename T, typename NS>
void KDTreeNode<T,NS>::print()
{
  std::cout << "idx =";
  for (int i = 0; i < this->n(); i++)
    std::cout << " " << this->idx_[this->idx(i)];
  std::cout << "\n";
  for (int i = 0; i < this->n_child(); i++)
    child(i)->print();
}

template<typename T, typename NS>
void KDTreeNode<T,NS>::nearest(const vec_type& p, KDTreeNode& Q, double &bound, int& minIdx)
{
  if (bound <= Q.nstats_.lowerBound(p)) return;               // check the bound

  if (Q.n_child() == 0) {                                     // reach a leaf
    for (int i = 0; i < Q.n(); i++) {                         // do naive search for
      double dist = arma::norm(p - Q.point(i), 2);            // minimum distance
      if (dist < bound) {
        bound = dist;
        minIdx = Q.idx(i);
      }
    }
  }
  else {
    for (int i = 0; i < Q.n_child(); i++) {                   // not a leaf
      KDTreeNode* child = Q.child(i);                         // recursive to children
      nearest(p, *child, bound, minIdx);
    }
  }
//  double min = DOUBLE_INFINITY;
//  for (int i = 0; i < Q.n(); i++) {
//    double dist = arma::norm(p-Q.point(i), 2);
//    if (dist < min) {
//      min_idx = Q.idx(i);
//      min = dist;
//    }
//  }
}

template<typename T, typename NS>
int KDTreeNode<T,NS>::nearest(const vec_type& p, KDTreeNode& Q)
{
  double minDist = DOUBLE_INFINITY;
  int minIdx = -1;
  nearest(p, Q, minDist, minIdx);
  return minIdx;
}

template<typename T, typename NS>
void KDTreeNode<T,NS>::allNearest(KDTreeNode& P, KDTreeNode& Q, std::vector<int>& minIdx, std::vector<double>& minDist)
{
  static int count = 0;
  static int c = 0;
  static int N = 1000000;
  if (P.upper_bound_ <= P.nstats_.boxLowerBound(Q.nstats_)) {
    count += P.n() * Q.n();
    if (count > c*N) {
//      std::cerr << " " << count;
      c = count / N + 1;
    }
    return;
  }
  if (P.n_child() == 0 && Q.n_child() == 0) {                // both are leaves
    P.upper_bound_ = 0;
    for (int i = 0; i < P.n(); i++) {                        // do naive search for min distance
      int pId = P.idx(i);
      for (int j = 0; j < Q.n(); j++) {
        int qId = Q.idx(j);
        double dist = arma::norm(P.point(i)-Q.point(j),2);
        if (minDist[pId] > dist) {
          minDist[pId] = dist;
          minIdx[pId] = qId;
        }
      }
      if (P.upper_bound_ < minDist[pId])                     // other ref node Q needs to be closer
        P.upper_bound_ = minDist[pId];                       // than this upper bound to be considered
    }
  }
  else if (P.n_child() == 0) {
    for (int j = 0; j < Q.n_child(); j++) {
      allNearest(P, *Q.child(j), minIdx, minDist);
      P.upper_bound_ = 0;
      for (int i = 0; i < P.n(); i++) {
        int pId = P.idx(i);
        if (P.upper_bound_ < minDist[pId])                     // other ref node Q needs to be closer
          P.upper_bound_ = minDist[pId];                       // than this upper bound to be considered
      }
    }
  }
  else if (Q.n_child() == 0) {
    P.upper_bound_ = 0;
    for (int i = 0; i < P.n_child(); i++) {
      allNearest(*P.child(i), Q, minIdx, minDist);
      double a = P.child(i)->upper_bound_;
      if (a > P.upper_bound_) P.upper_bound_ = a;
    }
  }
  else {
    P.upper_bound_ = 0;
    for (int i = 0; i < P.n_child(); i++) {
      for (int j = 0; j < Q.n_child(); j++)
        allNearest(*P.child(i), *Q.child(j), minIdx, minDist);
      double a = P.child(i)->upper_bound_;
      if (a > P.upper_bound_) P.upper_bound_ = a;
    }
  }
}

template<typename T, typename NS>
void KDTreeNode<T,NS>::setUpperBound(double u)
{
  upper_bound_ = u;
  for (int i = 0; i < this->n_child(); i++)
    this->child(i)->setUpperBound(u);
}

template<typename T, typename NS>
void KDTreeNode<T,NS>::allNearest(KDTreeNode& P, KDTreeNode& Q, std::vector<int>& minIdx)
{
  minIdx = std::vector<int>(P.n(), -1);
  std::vector<double> minDist(P.n(), DOUBLE_INFINITY);
  P.setUpperBound(DOUBLE_INFINITY);
  allNearest(P, Q, minIdx, minDist);
//  for (int i = 0; i < P.n(); i++) {
//    minIdx[P.idx(i)] = nearest(P.point(i), Q);
//  }
}

template<typename T, typename NS>
void KDTreeNode<T,NS>::combineStatistics()
{
//  std::cerr << "--combine n_child = " << this->n_child() << "\n";
  if (this->n_child() > 0) { // non-leaf, combine children stats
    std::vector<NS*> child_stats;
    for (int i = 0; i < this->n_child(); i++)
      child_stats.push_back(&(this->child(i)->nstats_));
    nstats_.combine(child_stats);
  }
  else { // leaf, build leaf stats
    for (int i = 0; i < this->n(); i++)
      nstats_.accumulate(this->point(i));
  }
}

// =============== Bounding Box ======================

template<typename T>
void BoundingBox<T>::combine(const std::vector<BoundingBox*> &child_boxes)
{
  if (child_boxes.size() == 0) return;
  minBox_ = child_boxes[0]->minBox_;
  maxBox_ = child_boxes[0]->maxBox_;
  int dim = minBox_.dim();
  assert(dim == maxBox_.dim());
  for (unsigned int i = 1; i < child_boxes.size(); i++) {
    assert(dim == child_boxes[i]->minBox_.dim() && dim == child_boxes[i]->maxBox_.dim());
    for (int d = 0; d < dim; d++) {
      double a = minBox_.get(d), b = child_boxes[i]->minBox_.get(d);
      if (a > b) minBox_.ref(d) = b;
      a = maxBox_.get(d); b = child_boxes[i]->maxBox_.get(d);
      if (a < b) maxBox_.ref(d) = b;
    }
    n_ += child_boxes[i]->n_;
  }
}

template<typename T>
void BoundingBox<T>::accumulate(const vec_type &p)
{
  int dim = p.dim();
  if (n_ == 0) {
    minBox_ = p; // first point
    maxBox_ = p;
    n_++;
  }
  else {                                     // later points
    for (int d = 0; d < dim; d++) {
      double a = minBox_.get(d), b = p.get(d);
      if (a > b) minBox_.ref(d) = b;
      a = maxBox_.get(d);
      if (a < b) maxBox_.ref(d) = b;
    }
    n_++;
  }
}

template<typename T>
double BoundingBox<T>::lowerBound(const vec_type &p)
{
  assert(p.dim() == minBox_.dim() && p.dim() == maxBox_.dim());
  double s = 0;
  int dim = p.dim();
  for (int i = 0; i < dim; i++) {
    double a = p.get(i);
    double b = minBox_.get(i);
    if (a < b) s += (b-a)*(b-a);
    else {
      b = maxBox_.get(i);
      if (a > b) s += (a-b)*(a-b);
    }
  }
  return ::sqrt(s);
}

template<typename T>
double BoundingBox<T>::boxLowerBound(const BoundingBox &box)
{
  int dim = box.minBox_.dim();
  assert(dim == box.maxBox_.dim() && dim == minBox_.dim() && dim == maxBox_.dim());
  double s = 0;
  for (int i = 0; i < dim; i++) {
    double a = minBox_.get(i), b = maxBox_.get(i);
    double c = box.minBox_.get(i), d = box.maxBox_.get(i);
    if (b < c) s += (c-b)*(c-b);
    else if (d < a) s += (a-d)*(a-d);
  }
  return ::sqrt(s);
}

template<typename T>
void BoundingBox<T>::print()
{
  std::cerr << "minBox_ =\n  " << minBox_ << "\n"
      << "maxBox_ = \n " << maxBox_ << "\n";
}

#endif /* KDTREE_IMPL_H_ */
